package com.cheems.cheemsspringaialibaba.Graph.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.cheems.cheemsspringaialibaba.Tools.SearchFromDify;
import org.springframework.ai.chat.client.ChatClient;

import java.util.Map;

import static com.cheems.cheemsspringaialibaba.Graph.config.HenauRoutingGraphConfiguration.availableRoutes;

/**
 * @author JTB
 */
public class RagNode implements NodeAction{


        private final ChatClient client;

        public RagNode(ChatClient client) {
            this.client = client;
        }

        @Override
        public Map<String, Object> apply(OverAllState state) {
            String inputText = (String) state.value("inputText").orElse("");
            String selectionLlm = (String) state.value("selectionLlm").orElse("");
            String result = client.prompt().system(availableRoutes.get(selectionLlm).toString())
                    .user(inputText)
                    .tools(new SearchFromDify())
                    .call().chatResponse().getResult().getOutput()
                    .getText();
            return Map.of("result", result);
        }

}
